import argparse
from concurrent import futures
import os
from pathlib import Path
from pysam import VariantFile, VariantHeader  # pylint: disable=E0611
from pre_base import PrepareBase, set_parser, logger


class PrepareSpliceAI(PrepareBase):
    input_file: list[Path]

    def header(self, chrom: str, size: int) -> VariantHeader:
        with VariantFile(str(self.input_file[0])) as reader:
            header = VariantHeader()
            header.contigs.add(id=chrom, length=size)
            for key, item in reader.header.info.items():
                header.info.add(id=key, type=item.type, number=item.number, description=item.description)
            return header

    def process_chrom(self, chrom: str, size: int, intervals: list[tuple[int, int]]):
        output_vcf = self.outdir / f"{chrom}.vcf"
        output_gz = output_vcf.with_suffix(".vcf.gz")
        if not output_gz.exists():
            with open(output_vcf, "w") as writer:
                writer.write(str(self.header(chrom, size)))
                for vcf_file in self.input_file:
                    with VariantFile(str(vcf_file)) as reader:
                        for start, end in intervals:
                            for row in reader.fetch(chrom.replace("chr", ""), start, end):
                                writer.write("chr" + str(row))
                            logger.info("processed %s:%s:%s %s", chrom, f"{start:,}", f"{end:,}", vcf_file.name)
            temp_dir = output_gz.parent / f"temp_{chrom}"
            temp_dir.mkdir(parents=True, exist_ok=True)
            command = f"bcftools sort {output_vcf} -T {temp_dir} -o {output_gz} -Oz"
            if os.system(command) > 0:
                raise RuntimeError(f"run command failed: {command}")
            command = f"bcftools index --tbi {output_gz}"
            if os.system(command) > 0:
                raise RuntimeError(f"run command failed: {command}")
            output_vcf.unlink(missing_ok=True)
        logger.info("processed %s", output_gz)

    def run(self):
        print(str(self.reference_file))
        self.outdir.mkdir(parents=True, exist_ok=True)
        with futures.ProcessPoolExecutor(max_workers=8) as executor:
            tasks = []
            for chrom, size, intervals in self.iter_reference():
                tasks.append(executor.submit(self.process_chrom, chrom, size, intervals))
            for task in futures.as_completed(tasks):
                task.result()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Prepare SpliceAI SNV")
    parser.add_argument("--input_file", "-i", nargs="+", type=Path, required=True, help="input, VCF files")
    set_parser(parser=parser)
    parser.set_defaults(func=lambda args: PrepareSpliceAI.model_validate(vars(args)).run())
    args = parser.parse_args()
    args.func(args)
